from torchvision import transforms as transforms
from torch.utils.data import DataLoader, Dataset
from skimage import io
from zisan.FileTools import getFiles
import os
import torch

class MnistDataset(Dataset):
    def __init__(self, root):
        self.root = root
        self.transforms = transforms.ToTensor() 
        paths = []
        labels = []
        for i in range(10):
            label = torch.zeros([10])
            label[i] = 1
            fs = getFiles(os.path.join(root,str(i)))
            ilabels = [label for _ in range(len(fs))]
            paths += fs
            labels += ilabels

        self.paths = paths
        self.labels = labels

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        img = io.imread(self.paths[idx],True) #as_gray =True
        img = self.transforms(img) #np.array ->tensor
        return img, self.labels[idx]

if __name__ == '__main__':
    ds = MnistDataset()
    dload=DataLoader(ds,shuffle=True,num_workers=0)
    for idx, (img, label) in enumerate(dload):
        print(img.shape)